package tech.aiflowy.ai.utils;

import com.agentsflex.core.prompt.template.TextPromptTemplate;
import tech.aiflowy.ai.entity.AiLlm;
import tech.aiflowy.ai.enums.TaskType;
import tech.aiflowy.ai.service.LLmTemplateService;
import tech.aiflowy.common.ai.ChatManager;
import tech.aiflowy.common.util.Maps;

/**
 * @author shadow
 * @llmClient 传入自己的llm实体,不传使用默认的
 */
public class TaskTypeClassifierUtil implements LLmTemplateService {
    private final AiLlm aiLlm;
    private final String promptTemplate;

    public TaskTypeClassifierUtil(AiLlm llmClient) {
        this.aiLlm = llmClient;
        // 初始化提示词模板
        this.promptTemplate = "请判断以下文本属于什么任务类型，仅返回任务类型（1/2/3/99），不要额外说明。文本：{{text}}";
    }

    @Override
    public TaskType format(String text) {
        //版本创建 PromptTemplate 并绑定变量
        TextPromptTemplate template = new TextPromptTemplate(promptTemplate);
        // 生成完整提示词
        String prompt = template.formatToString(Maps.of("text", text));
        // 调用 LLM
        String response;
        if (aiLlm == null) {
            response = ChatManager.getInstance().chat(text);
        } else {
            response = aiLlm.toLlm().chat(prompt);
        }
        // 解析结果
        return parseTaskType(response);
    }

    private TaskType parseTaskType(String response) {
        try {
            return TaskType.fromDes(response.trim());
        } catch (IllegalArgumentException e) {
            return TaskType.OTHER;
        }
    }
}
